# @Author：cnzdy
# @Email：cn_zdy@126.com
# @Time: 2021/10/5 10:33
# @File: dataset.py

import os
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
import torchvision.transforms as transforms
import torchvision
import matplotlib.pyplot as plt
import numpy as np
import torch

from pytorch.classifier.options import Options

normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])


def tsf_train(size, crop_size):
    return transforms.Compose(
        [
            transforms.Resize(size),
            transforms.RandomResizedCrop(crop_size),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize
        ])


def tsf_val(size):
    """ 验证集不用数据增强 """
    return transforms.Compose(
        [
            transforms.Resize(size),  # 尺寸规范
            transforms.CenterCrop(size),
            transforms.ToTensor(),  # 转化为tensor
            normalize
        ])


def tsf_test(size):
    """ 测试集不用数据增强"""
    return transforms.Compose(
        [
            transforms.Resize(size),  # 尺寸规范
            transforms.CenterCrop(size),
            transforms.ToTensor(),  # 转化为tensor
            normalize
        ])


def get_loader(path, transform=tsf_train(size=224, crop_size=224), shuffle=True):
    """使用 ImageFolder 载入指定文件夹下的数据集

    Example:
        dataloader = get_loader(path)
    :return: 加载器
    """
    dataset = []
    dataloader = None
    # print(f"batch_size: {batch_size}")
    if os.path.exists(path):
        # ImageFolder默认图像数据目录结构: root\dog\001.png, root\cat\001.png
        try:
            dataset = ImageFolder(path, transform=transform)
        except Exception:
            print("文件夹中没有子文件夹")
        dataloader = DataLoader(dataset, batch_size=Options.batch_size,
                                pin_memory=True, shuffle=shuffle,
                                num_workers=Options.num_workers)  # , collate_fn=fast_collate multiprocessing.cpu_count()
    return dataloader


def get_loaders(is_test=False):
    if is_test:
        Options.test_loader = get_loader(os.path.join(Options.root, 'test'),
                                         transform=tsf_test(size=224))
    else:
        Options.train_loader = get_loader(os.path.join(Options.root, 'train'))
        Options.val_loader = get_loader(os.path.join(Options.root, 'val'),
                                        transform=tsf_val(size=224))


def show_imgs(loader):
    """显示DataLoader中一个batch的所有图像"""
    dataiter = iter(loader)
    # images 是 Tensor类型
    images, _ = next(dataiter)
    # print(f"size: {images.size()}")

    img = torchvision.utils.make_grid(images)
    img = img / 2 + 0.5  # unnormalize
    npimg = img.numpy()

    # transpose()函数的作用就是调换数组的行列值的索引值，类似于求矩阵的转置
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


if __name__ == '__main__':
    # root = r'M:\classification\cat-dog\data\CatsDogsMini' + '/train'
    # dataloader = get_loader(root)
    # show_imgs(dataloader)

    #-------------------------------------------#
    get_loaders()
    show_imgs(Options.train_loader)

    # device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    # for i, data in enumerate(Options.train_loader, 0):
    #     input_batch = data[0].to(device, non_blocking=True)
    #     label_batch = data[1].to(device, non_blocking=True)
    #     print(f"input batch shape: {input_batch.shape}")
    #     print(f"label batch: {label_batch}")